package com.ljt.aicodehelp.ai.rag;

import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class RagConfig {

    @Bean
    public ContentRetriever contentRetriever(
            EmbeddingModel embeddingModel,
            EmbeddingStore<TextSegment> embeddingStore
    ) {
        // 准备好文档，准备好切片器
        var documentList = FileSystemDocumentLoader.loadDocuments("src/main/resources/docs");
        var documentByParagraphSplitter = new DocumentByParagraphSplitter(1000, 200);

        // 转换成向量存入数据库，
        // embedding应该翻译成转换向量。
        // 这里还可以用别的PGVector的ER数据库
        var ingestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(documentByParagraphSplitter)
                .textSegmentTransformer(
                        // 这里拼上文件名字，为了提高准确率。file_name硬编码就好
                        segment -> TextSegment.from(segment.metadata().getString("file_name")
                                +"\n"+ segment.text(),segment.metadata()))
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .build();
        ingestor.ingest(documentList);

        // 向量内容检索器
        // 最大返回5个结果
        // 自己应该自己调配实验出来的最小分数，匹配度分数小于0.75将被过滤掉

        // 放回一个向量内容检索器，就可以用上额外的知识库了，并且根据提问匹配回答了。
        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(5) // 最大返回5个结果
                .minScore(0.75) // 自己应该自己调配实验出来的最小分数，匹配度分数小于0.75将被过滤掉
                .build();
    }
}
